"""
可视化模块
负责数据探索和结果展示的可视化
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.gridspec import GridSpec


def setup_visualization_style():
    """
    设置可视化风格
    """
    # 设置Seaborn风格
    sns.set(style="whitegrid")

    # 设置中文字体支持（如果需要）
    try:
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
        plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    except:
        pass

    # 设置图表大小
    plt.rcParams['figure.figsize'] = (10, 6)

    print("可视化风格设置完成")


def plot_survival_distribution(df, title=None):
    """
    绘制生存状态分布

    参数:
        df (pd.DataFrame): 包含生存状态的数据框
        title (str): 图表标题

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 创建图表
    fig, ax = plt.subplots(figsize=(8, 6))

    # 计算生存状态计数
    survival_counts = df['Survived'].value_counts().sort_index()

    # 绘制条形图
    bars = ax.bar(['Not Survived', 'Survived'], survival_counts, color=['#ff6b6b', '#4ecdc4'])

    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width() / 2, height,
                f'{height} ({height/len(df)*100:.1f}%)',
                ha='center', va='bottom')

    # 设置标题和标签
    if title:
        ax.set_title(title)
    else:
        ax.set_title('Survival Distribution')

    ax.set_ylabel('Count')

    # 调整布局
    plt.tight_layout()

    return fig


def plot_survival_by_feature(df, feature, title=None, figsize=(10, 6)):
    """
    绘制按特征分组的生存率

    参数:
        df (pd.DataFrame): 包含生存状态和特征的数据框
        feature (str): 特征名称
        title (str): 图表标题
        figsize (tuple): 图表大小

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 创建图表
    fig, ax = plt.subplots(figsize=figsize)

    # 计算生存率
    survival_rate = df.groupby(feature)['Survived'].mean() * 100

    # 绘制条形图
    bars = ax.bar(survival_rate.index, survival_rate, color='#4ecdc4')

    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width() / 2, height,
                f'{height:.1f}%',
                ha='center', va='bottom')

    # 设置标题和标签
    if title:
        ax.set_title(title)
    else:
        ax.set_title(f'Survival Rate by {feature}')

    ax.set_ylabel('Survival Rate (%)')
    ax.set_xlabel(feature)

    # 调整布局
    plt.tight_layout()

    return fig


def plot_age_distribution_by_survival(df, title=None):
    """
    绘制按生存状态分组的年龄分布

    参数:
        df (pd.DataFrame): 包含年龄和生存状态的数据框
        title (str): 图表标题

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 创建图表
    fig, ax = plt.subplots(figsize=(10, 6))

    # 绘制直方图
    sns.histplot(data=df, x='Age', hue='Survived', kde=True, 
                 bins=30, alpha=0.6, ax=ax)

    # 设置标题和标签
    if title:
        ax.set_title(title)
    else:
        ax.set_title('Age Distribution by Survival')

    ax.set_xlabel('Age')
    ax.set_ylabel('Count')

    # 添加图例
    ax.legend(['Not Survived', 'Survived'])

    # 调整布局
    plt.tight_layout()

    return fig


def plot_fare_distribution_by_survival(df, title=None):
    """
    绘制按生存状态分组的票价分布

    参数:
        df (pd.DataFrame): 包含票价和生存状态的数据框
        title (str): 图表标题

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 创建图表
    fig, ax = plt.subplots(figsize=(10, 6))

    # 绘制直方图
    sns.histplot(data=df, x='Fare', hue='Survived', kde=True, 
                 bins=30, alpha=0.6, ax=ax)

    # 设置标题和标签
    if title:
        ax.set_title(title)
    else:
        ax.set_title('Fare Distribution by Survival')

    ax.set_xlabel('Fare')
    ax.set_ylabel('Count')

    # 添加图例
    ax.legend(['Not Survived', 'Survived'])

    # 调整布局
    plt.tight_layout()

    return fig


def plot_correlation_heatmap(df, title=None, figsize=(12, 10)):
    """
    绘制特征相关性热力图

    参数:
        df (pd.DataFrame): 数据框
        title (str): 图表标题
        figsize (tuple): 图表大小

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 计算相关性矩阵
    corr = df.corr()

    # 创建图表
    fig, ax = plt.subplots(figsize=figsize)

    # 绘制热力图
    mask = np.triu(np.ones_like(corr, dtype=bool))
    sns.heatmap(corr, mask=mask, annot=True, fmt='.2f', 
                cmap='coolwarm', vmin=-1, vmax=1, ax=ax)

    # 设置标题
    if title:
        ax.set_title(title)
    else:
        ax.set_title('Feature Correlation Heatmap')

    # 调整布局
    plt.tight_layout()

    return fig


def plot_feature_distributions(df, features, n_cols=3, figsize=(15, 10)):
    """
    绘制多个特征的分布图

    参数:
        df (pd.DataFrame): 数据框
        features (list): 特征名称列表
        n_cols (int): 每行显示的图表数量
        figsize (tuple): 图表大小

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 计算行数
    n_rows = (len(features) + n_cols - 1) // n_cols

    # 创建图表
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    axes = axes.flatten()

    # 绘制每个特征的分布
    for i, feature in enumerate(features):
        if feature in df.columns:
            ax = axes[i]

            # 根据特征类型选择不同的图表
            if df[feature].dtype == 'object' or len(df[feature].unique()) < 10:
                # 类别特征：绘制条形图
                sns.countplot(data=df, x=feature, ax=ax)
            else:
                # 数值特征：绘制直方图
                sns.histplot(data=df, x=feature, kde=True, ax=ax)

            # 设置标题
            ax.set_title(f'Distribution of {feature}')

            # 旋转x轴标签（如果需要）
            if len(df[feature].unique()) > 5:
                ax.tick_params(axis='x', rotation=45)

    # 隐藏多余的子图
    for i in range(len(features), len(axes)):
        axes[i].set_visible(False)

    # 调整布局
    plt.tight_layout()

    return fig


def plot_survival_by_multiple_features(df, features, n_cols=2, figsize=(15, 12)):
    """
    绘制多个特征与生存率的关系图

    参数:
        df (pd.DataFrame): 数据框
        features (list): 特征名称列表
        n_cols (int): 每行显示的图表数量
        figsize (tuple): 图表大小

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 计算行数
    n_rows = (len(features) + n_cols - 1) // n_cols

    # 创建图表
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    axes = axes.flatten()

    # 绘制每个特征与生存率的关系
    for i, feature in enumerate(features):
        if feature in df.columns:
            ax = axes[i]

            # 计算生存率
            survival_rate = df.groupby(feature)['Survived'].mean() * 100

            # 绘制条形图
            bars = ax.bar(survival_rate.index, survival_rate, color='#4ecdc4')

            # 添加数值标签
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width() / 2, height,
                        f'{height:.1f}%',
                        ha='center', va='bottom')

            # 设置标题
            ax.set_title(f'Survival Rate by {feature}')
            ax.set_ylabel('Survival Rate (%)')

            # 旋转x轴标签（如果需要）
            if len(survival_rate.index) > 5:
                ax.tick_params(axis='x', rotation=45)

    # 隐藏多余的子图
    for i in range(len(features), len(axes)):
        axes[i].set_visible(False)

    # 调整布局
    plt.tight_layout()

    return fig


def create_eda_dashboard(df, save_path=None):
    """
    创建探索性数据分析仪表板

    参数:
        df (pd.DataFrame): 数据框
        save_path (str, optional): 保存路径

    返回:
        matplotlib.figure.Figure: 生成的图表对象
    """
    # 创建图表
    fig = plt.figure(figsize=(20, 15))
    gs = GridSpec(3, 3, figure=fig)

    # 1. 生存状态分布
    ax1 = fig.add_subplot(gs[0, 0])
    survival_counts = df['Survived'].value_counts().sort_index()
    bars = ax1.bar(['Not Survived', 'Survived'], survival_counts, color=['#ff6b6b', '#4ecdc4'])
    for bar in bars:
        height = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width() / 2, height,
                f'{height} ({height/len(df)*100:.1f}%)',
                ha='center', va='bottom')
    ax1.set_title('Survival Distribution')
    ax1.set_ylabel('Count')

    # 2. 按性别分组的生存率
    ax2 = fig.add_subplot(gs[0, 1])
    survival_by_sex = df.groupby('Sex')['Survived'].mean() * 100
    bars = ax2.bar(survival_by_sex.index, survival_by_sex, color='#4ecdc4')
    for bar in bars:
        height = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width() / 2, height,
                f'{height:.1f}%',
                ha='center', va='bottom')
    ax2.set_title('Survival Rate by Sex')
    ax2.set_ylabel('Survival Rate (%)')

    # 3. 按船舱等级分组的生存率
    ax3 = fig.add_subplot(gs[0, 2])
    survival_by_pclass = df.groupby('Pclass')['Survived'].mean() * 100
    bars = ax3.bar(survival_by_pclass.index, survival_by_pclass, color='#4ecdc4')
    for bar in bars:
        height = bar.get_height()
        ax3.text(bar.get_x() + bar.get_width() / 2, height,
                f'{height:.1f}%',
                ha='center', va='bottom')
    ax3.set_title('Survival Rate by Pclass')
    ax3.set_ylabel('Survival Rate (%)')

    # 4. 年龄分布
    ax4 = fig.add_subplot(gs[1, 0])
    sns.histplot(data=df, x='Age', hue='Survived', kde=True, 
                 bins=30, alpha=0.6, ax=ax4)
    ax4.set_title('Age Distribution by Survival')
    ax4.set_xlabel('Age')
    ax4.set_ylabel('Count')

    # 5. 票价分布
    ax5 = fig.add_subplot(gs[1, 1])
    sns.histplot(data=df, x='Fare', hue='Survived', kde=True, 
                 bins=30, alpha=0.6, ax=ax5)
    ax5.set_title('Fare Distribution by Survival')
    ax5.set_xlabel('Fare')
    ax5.set_ylabel('Count')

    # 6. 按登船港口分组的生存率
    ax6 = fig.add_subplot(gs[1, 2])
    if 'Embarked' in df.columns:
        survival_by_embarked = df.groupby('Embarked')['Survived'].mean() * 100
        bars = ax6.bar(survival_by_embarked.index, survival_by_embarked, color='#4ecdc4')
        for bar in bars:
            height = bar.get_height()
            ax6.text(bar.get_x() + bar.get_width() / 2, height,
                    f'{height:.1f}%',
                    ha='center', va='bottom')
        ax6.set_title('Survival Rate by Embarked')
        ax6.set_ylabel('Survival Rate (%)')

    # 7. 家庭规模与生存率
    ax7 = fig.add_subplot(gs[2, 0])
    if 'FamilySize' in df.columns:
        survival_by_family = df.groupby('FamilySize')['Survived'].mean() * 100
        bars = ax7.bar(survival_by_family.index, survival_by_family, color='#4ecdc4')
        for bar in bars:
            height = bar.get_height()
            ax7.text(bar.get_x() + bar.get_width() / 2, height,
                    f'{height:.1f}%',
                    ha='center', va='bottom')
        ax7.set_title('Survival Rate by Family Size')
        ax7.set_ylabel('Survival Rate (%)')
        ax7.set_xlabel('Family Size')

    # 8. 头衔与生存率
    ax8 = fig.add_subplot(gs[2, 1])
    if 'Title' in df.columns:
        survival_by_title = df.groupby('Title')['Survived'].mean() * 100
        bars = ax8.bar(survival_by_title.index, survival_by_title, color='#4ecdc4')
        for bar in bars:
            height = bar.get_height()
            ax8.text(bar.get_x() + bar.get_width() / 2, height,
                    f'{height:.1f}%',
                    ha='center', va='bottom')
        ax8.set_title('Survival Rate by Title')
        ax8.set_ylabel('Survival Rate (%)')
        ax8.tick_params(axis='x', rotation=45)

    # 9. 特征相关性热力图
    ax9 = fig.add_subplot(gs[2, 2])
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    corr = df[numeric_cols].corr()
    mask = np.triu(np.ones_like(corr, dtype=bool))
    sns.heatmap(corr, mask=mask, annot=True, fmt='.2f', 
                cmap='coolwarm', vmin=-1, vmax=1, ax=ax9)
    ax9.set_title('Feature Correlation Heatmap')

    # 调整布局
    plt.tight_layout()

    # 保存图表
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"EDA仪表板已保存到: {save_path}")

    return fig
